package container

import (
	"strconv"
	"strings"

	"k8s.io/apimachinery/pkg/runtime/schema"

	"k8s.io/client-go/discovery"
	// memory "k8s.io/client-go/discovery/cached"
	"context"

	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/restmapper"

	// "os"
	//	"bufio"
	"flag"
	"fmt"
	"runtime"
	"sugar/utils"

	//	"os"
	"path/filepath"

	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime/serializer/yaml"

	"k8s.io/client-go/discovery/cached/memory"
	"k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"

	// clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
	"k8s.io/client-go/util/homedir"
	// unver "k8s.io/kubernetes/pkg/api/unversioned"
	//	"k8s.io/client-go/util/retry"
)

const (
	ResourceRequestsGPU apiv1.ResourceName = "requests.nvidia.com/gpu"
	ResourceLimitsGPU   apiv1.ResourceName = "limits.nvidia.com/gpu"
	ResourceGPU         apiv1.ResourceName = "nvidia.com/gpu"
)

type ComputerSource struct {
	CPU_Allocable int64 `json:"cpu_allocable"`
	CPU_Total     int64 `json:"cpu_total"`
	CPU_Per_Max   int64 `json:"cpu_per_max"`
	MEN_Allocable int64 `json:"men_allocable"`
	MEN_Total     int64 `json:"men_total"`
	MEN_Per_Max   int64 `json:"men_per_max"`
	GPU_Allocable int64 `json:"gpu_allocable"`
	GPU_Total     int64 `json:"gpu_total"`
	GPU_Per_Max   int64 `json:"gpu_per_max"`
}

// type K8s_operation interface {
// 	CreateTrainjob(trainjob_id string, gpu_count int) error
// 	DeleteTrainjob(trainjob_id string) error
// 	//CreatePredictjob(predictjob_id string) (string, error)
// 	//DeletePredictjob(predictjob_id string) error
// 	CreateApp(app_id string, host_port int) error
// 	DeleteApp(app_id string) error
// 	GetComputerSource() (*ComputerSource, error)

// 	GetTrainjobDeployStatus(trainjob_id string) (string, error)
// 	GetAppDeployStatus(app_id string) (string, error)
// 	GetDeployStatus(deploy_name string) (string, error)
// }

type KubernetesOperation struct {
	clientset       *kubernetes.Clientset
	dynamicClient   dynamic.Interface
	discoveryClient *discovery.DiscoveryClient

	mappingAppServer *meta.RESTMapping
	mappingTrainJob  *meta.RESTMapping
}

func NewK8sOp() *KubernetesOperation {
	var config *restclient.Config
	var err error
	if utils.GetEnvDefault("RUN_MODE", "HOST") == "HOST" {
		var kubeconfig string
		if home := homedir.HomeDir(); home != "" {
			kubeconfig = filepath.Join(home, ".kube", "config") //flag.String("kubeconfig", filepath.Join(home, ".kube", "config"), "(optional) absolute path to the kubeconfig file")
		} else {
			kubeconfig = "" //flag.String("kubeconfig", "", "absolute path to the kubeconfig file")
		}
		flag.Parse()

		config, err = clientcmd.BuildConfigFromFlags("", kubeconfig)
	} else {
		config, err = restclient.InClusterConfig()
		// k8s_url := fmt.Sprintf("https://%s:%s", os.Getenv("KUBERNETES_SERVICE_HOST"), os.Getenv("KUBERNETES_SERVICE_PORT"))
		// config = new(restclient.Config)
		// config.BearerTokenFile = "/var/run/secrets/kubernetes.io/serviceaccount/token"
		// config.Host = k8s_url
		// config.CAFile = ""
		//config.CAFile = nil
	}
	if err != nil {
		return nil
	}

	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		return nil
	} else {
		fmt.Println("connect k8s success")
	}

	dynamicClient, err := dynamic.NewForConfig(config)
	if err != nil {
		return nil
	} else {
		fmt.Println("connect k8s success")
	}

	dc, err := discovery.NewDiscoveryClientForConfig(config)
	if err != nil {
		return nil
	}

	a := new(KubernetesOperation)
	if err != nil {
		return nil
	}
	a.clientset = clientset
	a.dynamicClient = dynamicClient
	a.discoveryClient = dc
	// a := &k8s_operation_s{clientset: clientset}
	gvkAppServer := schema.GroupVersionKind{Group: "sugar.sugar.com", Version: "v1", Kind: "AppServer"}
	gvkTrainJob := schema.GroupVersionKind{Group: "sugar.sugar.com", Version: "v1", Kind: "DistributeTrainJob"}

	mapper := restmapper.NewDeferredDiscoveryRESTMapper(memory.NewMemCacheClient(a.discoveryClient))

	a.mappingAppServer, err = mapper.RESTMapping(gvkAppServer.GroupKind(), gvkAppServer.Version)
	if err != nil {
		return nil
	}
	a.mappingTrainJob, err = mapper.RESTMapping(gvkTrainJob.GroupKind(), gvkTrainJob.Version)
	if err != nil {
		return nil
	}
	runtime.SetFinalizer(a, DelK8sOp)
	return a
}

func DelK8sOp(c *KubernetesOperation) {
	// fmt.Printf("db close!\n")
	// if c.clientset != nil {
	// 	c.clientset..Close()
	// }
	//time.Sleep(time.Second)
}

func (c *KubernetesOperation) GetComputerSource() (*ComputerSource, error) {
	nodes, err := c.clientset.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, err
	}
	var cpuTotal int64 = 0
	var menTotal int64 = 0
	var gpuTotal int64 = 0
	var cpuAllocTotal int64 = 0
	var menAllocTotal int64 = 0
	var gpuAllocTotal int64 = 0
	var cpuPerMax int64 = 0
	var menPerMax int64 = 0
	var gpuPerMax int64 = 0

	for _, nds := range nodes.Items {
		fmt.Printf("NodeName: %s\n", nds.Name)
		if cpuCount, ok := nds.Status.Capacity.Cpu().AsInt64(); ok {
			cpuTotal += cpuCount
		}
		if cpuAlloc, ok := nds.Status.Allocatable.Cpu().AsInt64(); ok {
			cpuAllocTotal += cpuAlloc
			if cpuAlloc > cpuPerMax {
				cpuPerMax = cpuAlloc
			}
		}
		if menCount, ok := nds.Status.Capacity.Memory().AsInt64(); ok {
			menTotal += menCount
		}
		if menAlloc, ok := nds.Status.Allocatable.Memory().AsInt64(); ok {
			menAllocTotal += menAlloc
			if menAlloc > menPerMax {
				menPerMax = menAlloc
			}
		}

		gpuResource := nds.Status.Capacity[ResourceGPU]
		//var capacity apiv1.ResourceList = nds.Status.Capacity.(map[apiv1.ResourceName]apiv1.resource.Quantity)
		if gpuCount, ok := gpuResource.AsInt64(); ok {
			gpuTotal += gpuCount
		}
		//apiv1.ResourceList
		gpuAllocResource := nds.Status.Allocatable[ResourceGPU]
		if gpuAllocCount, ok := gpuAllocResource.AsInt64(); ok {
			gpuAllocTotal += gpuAllocCount
			if gpuAllocCount > gpuPerMax {
				gpuPerMax = gpuAllocCount
			}
		}
	}
	res := &ComputerSource{
		CPU_Total:     cpuTotal,
		MEN_Total:     menTotal,
		GPU_Total:     gpuTotal,
		CPU_Allocable: cpuAllocTotal,
		MEN_Allocable: menAllocTotal,
		GPU_Allocable: gpuAllocTotal,
		CPU_Per_Max:   cpuPerMax,
		MEN_Per_Max:   menPerMax,
		GPU_Per_Max:   gpuPerMax,
	}
	return res, nil
}

// func (c *KubernetesOperation) GetDeployStatus(deploy_name string) (string, error) {
// 	deploymentsClient := c.clientset.AppsV1().Deployments("sugar")
// 	deploy, err := deploymentsClient.Get(deploy_name, metav1.GetOptions{})
// 	if err != nil {
// 		return "", err
// 	}
// 	podSelector, err := metav1.LabelSelectorAsSelector(deploy.Spec.Selector)
// 	if err != nil {
// 		return "", err
// 	}
// 	podClient := c.clientset.CoreV1().Pods("sugar")
// 	podList, err := podClient.List(metav1.ListOptions{LabelSelector: podSelector.String()})
// 	if err != nil {
// 		return "", err
// 	}
// 	// fmt.Println(podList.Items[0].Name)
// 	// fmt.Println(podList.Items[0].Status.Phase)
// 	return string(podList.Items[0].Status.Phase), nil
// }

// func (c *KubernetesOperation) GetTrainjobDeployStatus(trainjob_id string) (string, error) {
// 	container_name := fmt.Sprintf("trainjob-%s", trainjob_id)
// 	return c.GetDeployStatus(container_name)
// }

// func (c *KubernetesOperation) GetAppDeployStatus(app_id string) (string, error) {
// 	container_name := fmt.Sprintf("app-%s", app_id)
// 	return c.GetDeployStatus(container_name)
// }

const appYamlTemplate = `
apiVersion: sugar.sugar.com/v1
kind: AppServer
metadata:
  name: NAME
spec:
  # Add fields here
  replicasmin: REPLICAS_MIN
  replicasmax: REPLICAS_MAX
  modelfolder: MODEL_FOLDER
  modelfile: "MODEL_FILE"
  modelparam: "MODEL_PARAM"
  modelclass: "MODEL_CLASS"
  port:
    targetPort: 9002
    nodePort: PORT
    port: PORT
  modelvolume:
    name: "modelfolder"
	persistentVolumeClaim:
      claimName: sugar-model-pvc
  modelparamvolume:
    name: "modelparamfolder"
	persistentVolumeClaim:
      claimName: sugar-model-params-pvc
  metrics:
    - type: Resource
      resource:
        name: cpu
        target:
          type: Utilization
          averageUtilization: 50
  resources:
    requests:
      cpu: CPU_REQUEST
	  memory: MEMORY_REQUEST
  statusinformhook: STATUS_INFORM_HOOK
`

//
func (c *KubernetesOperation) CreateApp(appID string, replicaMin int, replicaMax int, modelFile string, modelParam string, modelClass string,
	servicePort int, cpuRequest string, memoryRequest string, modelFolder string, statusHookURL string) error {
	var appYaml string = appYamlTemplate
	appYaml = strings.ReplaceAll(appYaml, "NAME", "App-"+appID)
	appYaml = strings.ReplaceAll(appYaml, "REPLICAS_MIN", strconv.Itoa(replicaMin))
	appYaml = strings.ReplaceAll(appYaml, "REPLICAS_MAX", strconv.Itoa(replicaMax))
	appYaml = strings.ReplaceAll(appYaml, "MODEL_FOLDER", modelFolder)
	appYaml = strings.ReplaceAll(appYaml, "MODEL_FILE", modelFile)
	appYaml = strings.ReplaceAll(appYaml, "MODEL_PARAM", modelParam)
	appYaml = strings.ReplaceAll(appYaml, "MODEL_CLASS", modelClass)

	// nfsServer := utils.GetEnvDefault("NFS_SERVER", "192.168.2.89")
	// appYaml = strings.ReplaceAll(appYaml, "NFS_SERVER", nfsServer)

	// appYaml = strings.ReplaceAll(appYaml, "MODEL_VOLUME_PATH", modelVolumePath)
	// appYaml = strings.ReplaceAll(appYaml, "MODEL_PARAM_VOLUME_PATH", modelParamVolumePath)
	appYaml = strings.ReplaceAll(appYaml, "CPU_REQUEST", cpuRequest)
	appYaml = strings.ReplaceAll(appYaml, "MEMORY_REQUEST", memoryRequest)

	appYaml = strings.ReplaceAll(appYaml, "PORT", strconv.Itoa(servicePort))
	appYaml = strings.ReplaceAll(appYaml, "STATUS_INFORM_HOOK", statusHookURL)
	obj := &unstructured.Unstructured{}

	// decode YAML into unstructured.Unstructured
	dec := yaml.NewDecodingSerializer(unstructured.UnstructuredJSONScheme)
	_, _, err := dec.Decode([]byte(appYaml), nil, obj)
	if err != nil {
		return err
	}
	// unstructured.
	_, err = c.dynamicClient.Resource(c.mappingAppServer.Resource).Namespace("sugar").Create(context.TODO(), obj, metav1.CreateOptions{})
	return err
}

//
func (c *KubernetesOperation) DeleteApp(appID string) error {
	err := c.dynamicClient.Resource(c.mappingAppServer.Resource).Namespace("sugar").Delete(context.TODO(), "App-"+appID, metav1.DeleteOptions{})
	return err
}

const trainJobYamlTemplate = `
apiVersion: sugar.sugar.com/v1
kind: DistributeTrainJob
metadata:
  name: NAME
spec:
  replicas: 1
  selector:
    matchLabels:
      name: NAME
  ports:
  - port: 90
    name: test
  modelvolume:
    name: modelfolder
    persistentVolumeClaim:
      claimName: sugar-model-pvc
  modelparamvolume:
    name: modelparamfolder
    persistentVolumeClaim:
      claimName: sugar-model-param-pvc
  datasetvolume:
    name: datasetfolder
    persistentVolumeClaim:
      claimName: sugar-dataset-pvc
  logvolume:
    name: logfolder
    persistentVolumeClaim:
      claimName: sugar-log-pvc
  modelsavepath: MODEL_SAVE_PATH
  modelfolder: MODEL_FOLDER
  modelfile: MODEL_FILE
  modelclass: MODEL_CLASS
  modelcheckpoint: ""
  modelparams: "MODEL_PARAMS"
  trainparams: "TRAIN_PARAMS"
  traindataset: TRAIN_DATASET
  validatedataset: VALIDATE_DATASET
  destscore: "DEST_SCORE"
  maxtrials: MAX_TRAILS
  useautoml: USE_AUTO_ML
  statusinformhook: STATUS_INFORM_HOOK
  resources:
    requests:
      cpu: CPU_REQUEST
      memory: MEMORY_REQUEST
`

//
func (c *KubernetesOperation) CreateTrainjob(trainJobID string, modelSavePath string, modelFile string, modelClass string,
	modelParams string, trainParams string, trainDataset string, validateDataset string, destScore string,
	maxTrials int, useAutoMl int, cpuRequest string, memoryRequest string, gpuCount int, modelFolder string,
	statusHookURL string) error {
	modelParams = strings.ReplaceAll(modelParams, "\"", "\\\"")
	trainParams = strings.ReplaceAll(trainParams, "\"", "\\\"")
	var trainJobYaml string = trainJobYamlTemplate
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "NAME", "trainjob-"+trainJobID)
	// nfsServer := utils.GetEnvDefault("NFS_SERVER", "192.168.2.89")
	// trainJobYaml = strings.ReplaceAll(trainJobYaml, "NFS_SERVER", nfsServer)
	// trainJobYaml = strings.ReplaceAll(trainJobYaml, "MODEL_FOLDER_PATH", modelFolderPath)
	// trainJobYaml = strings.ReplaceAll(trainJobYaml, "MODEL_PARAM_FOLDER_PATH", modelParamFolderPath)
	// trainJobYaml = strings.ReplaceAll(trainJobYaml, "DATASET_FOLDER_PATH", datasetFolderPath)
	// trainJobYaml = strings.ReplaceAll(trainJobYaml, "LOG_FOLDER_PATH", logFolderPath)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "MODEL_SAVE_PATH", modelSavePath)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "MODEL_FOLDER", modelFolder)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "MODEL_FILE", modelFile)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "MODEL_CLASS", modelClass)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "MODEL_PARAMS", modelParams)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "TRAIN_PARAMS", trainParams)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "TRAIN_DATASET", trainDataset)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "VALIDATE_DATASET", validateDataset)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "DEST_SCORE", destScore)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "MAX_TRAILS", strconv.Itoa(maxTrials))
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "USE_AUTO_ML", strconv.Itoa(useAutoMl))
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "CPU_REQUEST", cpuRequest)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "MEMORY_REQUEST", memoryRequest)
	trainJobYaml = strings.ReplaceAll(trainJobYaml, "STATUS_INFORM_HOOK", statusHookURL)

	obj := &unstructured.Unstructured{}

	// decode YAML into unstructured.Unstructured
	dec := yaml.NewDecodingSerializer(unstructured.UnstructuredJSONScheme)
	_, _, err := dec.Decode([]byte(trainJobYaml), nil, obj)
	if err != nil {
		return err
	}
	// unstructured.
	_, err = c.dynamicClient.Resource(c.mappingTrainJob.Resource).Namespace("sugar").Create(context.TODO(), obj, metav1.CreateOptions{})
	return err
}

func (c *KubernetesOperation) GetTrainjobStatus(trainJobID string) string {
	return "RUNNING"
}

//
func (c *KubernetesOperation) DeleteTrainjob(trainJobID string) error {
	err := c.dynamicClient.Resource(c.mappingTrainJob.Resource).Namespace("sugar").Delete(context.TODO(), "TrainJob-"+trainJobID, metav1.DeleteOptions{})
	return err
}

func int32Ptr(i int32) *int32 { return &i }
